{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of online analysis using OnACID\n",
    "\n",
    "Complete pipeline for online processing using CaImAn Online (OnACID).\n",
    "The demo demonstrates the analysis of a sequence of files using the CaImAn online\n",
    "algorithm. The steps include i) motion correction, ii) tracking current \n",
    "components, iii) detecting new components, iv) updating of spatial footprints.\n",
    "The script demonstrates how to construct and use the params and online_cnmf\n",
    "objects required for the analysis, and presents the various parameters that\n",
    "can be passed as options. A plot of the processing time for the various steps\n",
    "of the algorithm is also included.\n",
    "@author: Eftychios Pnevmatikakis @epnev\n",
    "Special thanks to Andreas Tolias and his lab at Baylor College of Medicine\n",
    "for sharing the data used in this demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython import get_ipython\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import caiman as cm\n",
    "from caiman.source_extraction import cnmf\n",
    "from caiman.utils.utils import download_demo\n",
    "\n",
    "logfile = None # Replace with a path if you want to log to a file\n",
    "logger = logging.getLogger('caiman')\n",
    "# Set to logging.INFO if you want much output, potentially much more output\n",
    "logger.setLevel(logging.WARNING)\n",
    "logfmt = logging.Formatter('%(relativeCreated)12d [%(filename)s:%(funcName)20s():%(lineno)s] [%(process)d] %(message)s')\n",
    "if logfile is not None:\n",
    "    handler = logging.FileHandler(logfile)\n",
    "else:\n",
    "    handler = logging.StreamHandler()\n",
    "handler.setFormatter(logfmt)\n",
    "logger.addHandler(handler)\n",
    "\n",
    "try:\n",
    "    if __IPYTHON__:\n",
    "        get_ipython().run_line_magic('load_ext', 'autoreload')\n",
    "        get_ipython().run_line_magic('autoreload', '2')\n",
    "except NameError:\n",
    "    pass\n",
    "\n",
    "import bokeh.plotting as bpl\n",
    "bpl.output_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First download the data\n",
    "\n",
    "The function ```download_demo``` will look for the datasets ```Tolias_mesoscope_*.hdf5``` ins your caiman_data folder inside the subfolder specified by the variable ```fld_name``` and will download the files if they do not exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fld_name = 'Mesoscope'                              # folder inside ./example_movies where files will be saved\n",
    "fnames = []\n",
    "fnames.append(download_demo('Tolias_mesoscope_1.hdf5',fld_name))\n",
    "fnames.append(download_demo('Tolias_mesoscope_2.hdf5',fld_name))\n",
    "fnames.append(download_demo('Tolias_mesoscope_3.hdf5',fld_name))\n",
    "\n",
    "print(fnames)                                          # your list of files should look something like this"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up some parameters\n",
    "\n",
    "Here we set up some parameters for running OnACID. We use the same `params` object as in batch processing with CNMF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fr = 15                                                             # frame rate (Hz)\n",
    "decay_time = 0.5                                                    # approximate length of transient event in seconds\n",
    "gSig = (4,4)                                                        # expected half size of neurons\n",
    "p = 1                                                               # order of AR indicator dynamics\n",
    "min_SNR = 1                                                         # minimum SNR for accepting new components\n",
    "rval_thr = 0.90                                                     # correlation threshold for new component inclusion\n",
    "ds_factor = 1                                                       # spatial downsampling factor (increases speed but may lose some fine structure)\n",
    "gnb = 2                                                             # number of background components\n",
    "gSig = tuple(np.ceil(np.array(gSig)/ds_factor).astype('int'))       # recompute gSig if downsampling is involved\n",
    "mot_corr = True                                                     # flag for online motion correction \n",
    "pw_rigid = False                                                    # flag for pw-rigid motion correction (slower but potentially more accurate)\n",
    "max_shifts_online = np.ceil(10./ds_factor).astype('int')            # maximum allowed shift during motion correction\n",
    "sniper_mode = True                                                  # flag using a CNN to detect new neurons (o/w space correlation is used)\n",
    "init_batch = 200                                                    # number of frames for initialization (presumably from the first file)\n",
    "expected_comps = 500                                                # maximum number of expected components used for memory pre-allocation (exaggerate here)\n",
    "dist_shape_update = True                                            # flag for updating shapes in a distributed way\n",
    "min_num_trial = 10                                                  # number of candidate components per frame     \n",
    "K = 2                                                               # initial number of components\n",
    "epochs = 2                                                          # number of passes over the data\n",
    "show_movie = False                                                  # show the movie with the results as the data gets processed\n",
    "\n",
    "params_dict = {'fnames': fnames,\n",
    "               'fr': fr,\n",
    "               'decay_time': decay_time,\n",
    "               'gSig': gSig,\n",
    "               'p': p,\n",
    "               'min_SNR': min_SNR,\n",
    "               'rval_thr': rval_thr,\n",
    "               'ds_factor': ds_factor,\n",
    "               'nb': gnb,\n",
    "               'motion_correct': mot_corr,\n",
    "               'init_batch': init_batch,\n",
    "               'init_method': 'bare',\n",
    "               'normalize': True,\n",
    "               'expected_comps': expected_comps,\n",
    "               'sniper_mode': sniper_mode,\n",
    "               'dist_shape_update' : dist_shape_update,\n",
    "               'min_num_trial': min_num_trial,\n",
    "               'K': K,\n",
    "               'epochs': epochs,\n",
    "               'max_shifts_online': max_shifts_online,\n",
    "               'pw_rigid': pw_rigid,\n",
    "               'show_movie': show_movie}\n",
    "opts = cnmf.params.CNMFParams(params_dict=params_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now run the CaImAn online algorithm (OnACID).\n",
    "\n",
    "The first ```initbatch``` frames are used for initialization purposes. The initialization method chosen here `bare` will only search for a small number of neurons and is mostly used to initialize the background components. Initialization with the full CNMF can also be used by choosing `cnmf`.\n",
    "\n",
    "We first create an `OnACID` object located in the module `online_cnmf` and we pass the parameters similarly to the case of batch processing. We then run the algorithm using the `fit_online` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm = cnmf.online_cnmf.OnACID(params=opts)\n",
    "cnm.fit_online()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optionally save results and do some plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info('Number of components: ' + str(cnm.estimates.A.shape[-1]))\n",
    "Cn = cm.load(fnames[0], subindices=slice(0,500)).local_correlations(swap_dim=False)\n",
    "cnm.estimates.plot_contours(img=Cn)\n",
    "# cnm.save('demo_OnACID_mesoscope_results.hdf5') # Or add /full/path/to/file.hdf5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## View components\n",
    "\n",
    "Now inspect the components extracted by OnACID. Note that if single pass was used then several components would be non-zero only for the part of the time interval indicating that they were detected online by OnACID.\n",
    "\n",
    "Note that if you get data rate error you can start Jupyter notebooks using:\n",
    "'jupyter lab --ZMQChannelsWebsocketConnection.iopub_data_rate_limit=1.0e10'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm.estimates.nb_view_components(img=Cn, denoised_color='red');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot timing\n",
    "The plot below shows the time spent on each part of the algorithm (motion correction, tracking of current components, detect new components, update shapes) for each frame. Note that if you displayed a movie while processing the data (`show_movie=True`) the time required to generate this movie will be included here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_motion = 1e3*np.array(cnm.t_motion)\n",
    "T_detect = 1e3*np.array(cnm.t_detect)\n",
    "T_shapes = 1e3*np.array(cnm.t_shapes)\n",
    "T_online = 1e3*np.array(cnm.t_online) - T_motion - T_detect - T_shapes\n",
    "plt.figure()\n",
    "plt.stackplot(np.arange(len(T_motion)), T_motion, T_online, T_detect, T_shapes)\n",
    "plt.legend(labels=['motion', 'process', 'detect', 'shapes'], loc=2)\n",
    "plt.title('Processing time allocation')\n",
    "plt.xlabel('Frame #')\n",
    "plt.ylabel('Processing time [ms]')\n",
    "plt.ylim([0,140])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "caiman_pytorch2",
   "language": "python",
   "name": "caiman_pytorch2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
